from ollama import Client
import tkinter as tk
from tkinter import scrolledtext

# 从 ./api-key.txt 读取TOKEN
with open("./api-key.txt", "r") as f:
    TOKEN = f.read().strip()


# 获取 Ollama 客户端
def get_client(token: str = None):
    client = Client(
        host="https://myoi.ynuosa.org/ollama/",
        headers={"Authorization": f"Bearer {token}"},
    )
    return client


# 聊天接口：向 Ollama 发送消息
def chat_with_ollama(client, messages):
    try:
        response = client.chat(model="qwen2.5:14b", messages=messages)
        return response['message']['content']
    except Exception as e:
        return f"接口调用出错：{e}"


# 金额转换为中文大写
def num_to_chinese(num):
    units = ["", "拾", "佰", "仟"]
    big_units = ["", "万", "亿", "兆"]
    digits = "零壹贰叁肆伍陆柒捌玖"
    integer_part, decimal_part = str(num).split('.')
    
    integer_part = integer_part[::-1]
    chinese_integer = []
    zero_count = 0
    
    for i in range(0, len(integer_part), 4):
        segment = integer_part[i:i+4][::-1]
        segment_chinese = []
        for j, digit in enumerate(segment):
            if digit != '0':
                if zero_count > 0:
                    segment_chinese.append(digits[0])
                    zero_count = 0
                segment_chinese.append(digits[int(digit)])
                segment_chinese.append(units[j])
            else:
                zero_count += 1
        if segment_chinese:
            segment_chinese = segment_chinese[:-1]  # Remove the last unit
            segment_chinese.append(big_units[len(segment_chinese) // 4])
            chinese_integer = segment_chinese + chinese_integer
        elif zero_count < 4:  # Add zero for segments with all zeros but not at the end
            chinese_integer = [digits[0]] + chinese_integer
            zero_count = 0
    
    chinese_integer = [c for c in chinese_integer if c != digits[0] or chinese_integer.index(c) == len(chinese_integer) - 1]  # Remove leading zeros
    chinese_integer = ''.join(chinese_integer)
    
    # Process decimal part
    chinese_decimal = []
    jiao = int(decimal_part[0])
    fen = int(decimal_part[1])
    if jiao:
        chinese_decimal.append(digits[jiao] + "角")
    if fen:
        chinese_decimal.append(digits[fen] + "分")
    chinese_decimal = ''.join(chinese_decimal) if chinese_decimal else "整"
    
    return chinese_integer + "元" + chinese_decimal


# GUI 界面实现
class QAApp:
    def __init__(self, root, client):
        self.client = client
        self.root = root
        self.root.title("智能问答助手")

        # 输入框标签
        self.input_label = tk.Label(root, text="请输入你的问题：", font=("Arial", 12))
        self.input_label.pack(pady=5)

        # 输入框
        self.input_box = tk.Entry(root, width=60, font=("Arial", 12))
        self.input_box.pack(pady=5)

        # 按钮区
        self.button_frame = tk.Frame(root)
        self.button_frame.pack(pady=5)

        self.joke_button = tk.Button(self.button_frame, text="趣味问答", command=lambda: self.get_answer("joke"))
        self.joke_button.grid(row=0, column=0, padx=5)

        self.story_button = tk.Button(self.button_frame, text="小故事生成", command=lambda: self.get_answer("story"))
        self.story_button.grid(row=0, column=1, padx=5)

        self.quote_button = tk.Button(self.button_frame, text="每日心灵鸡汤", command=lambda: self.get_answer("quote"))
        self.quote_button.grid(row=0, column=2, padx=5)

        self.translate_to_classical_chinese = tk.Button(
            self.button_frame, text="翻译成文言文", command=lambda: self.get_answer("translate_to_classical_chinese")
        )
        self.translate_to_classical_chinese.grid(row=1, column=0, padx=5)

        self.couplet_button = tk.Button(self.button_frame, text="对联生成", command=lambda: self.get_answer("couplet"))
        self.couplet_button.grid(row=0, column=3, padx=5)

        self.money_to_chinese_button = tk.Button(
            self.button_frame, text="金额转换为中文大写", command=self.convert_money_to_chinese
        )
        self.money_to_chinese_button.grid(row=1, column=3, padx=5)

        # 输出框标签
        self.output_label = tk.Label(root, text="回答：", font=("Arial", 12))
        self.output_label.pack(pady=5)

        # 输出文本框
        self.output_box = scrolledtext.ScrolledText(root, width=60, height=10, font=("Arial", 12))
        self.output_box.pack(pady=5)

    def get_answer(self, mode):
        """根据模式生成不同的 Prompt 并调用 Ollama"""
        user_input = self.input_box.get()

        if mode == "joke":
            prompt = f"你现在是一个幽默大师，请用轻松、有趣的调度回答这个问题：{user_input}"
        elif mode == "story":
            prompt = f"请为我写一个简短的小故事，故事的主题是：{user_input}，字数在100-200字之间，生动有趣。"
        elif mode == "quote":
            prompt = "请生成一句充满正能量的励志句子，鼓励读者开始美好的一天。"
        elif mode == "translate_to_classical_chinese":
            prompt = f"请把这段文字：'{user_input}'翻译成文言文。在你的回答中，你只输出翻译得到的文言文版本。你要给出2~3个不同的翻译版本。"
        elif mode == "couplet":
            prompt = f"请根据输入的关键词：'{user_input}'，创作一副对联。要求：上联、下联、檀题均简洁、合谏，具有美感和情意。只输出对联内容，不要输出其他无关的东西"
        else:
            prompt = user_input

        # 调用 Ollama API
        messages = [{"role": "user", "content": prompt}]
        response = chat_with_ollama(self.client, messages)

        # 显示输出结果
        self.output_box.delete("1.0", tk.END)
        self.output_box.insert(tk.END, response)

    def convert_money_to_chinese(self):
        """将输入金额转换为中文大写"""
        try:
            amount = float(self.input_box.get())
            chinese_amount = num_to_chinese(amount)
            self.output_box.delete("1.0", tk.END)
            self.output_box.insert(tk.END, chinese_amount)
        except ValueError:
            self.output_box.delete("1.0", tk.END)
            self.output_box.insert(tk.END, "请输入有效的数字金额")


# 主程序
if __name__ == "__main__":
    client = get_client(TOKEN)  # 初始化 Ollama 客户端
    root = tk.Tk()
    app = QAApp(root, client)